# -*- coding: utf-8 -*-

import cv2
import tensorflow as tf
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
from model import CNN

# 定义常量
BASE_DIR = '/Users/patrick/Workspace/git/aphf_score_recognize/'

class Predict(object):
    def __init__(self):
        latest = tf.train.latest_checkpoint(BASE_DIR + 'model/ckpt')
        self.cnn = CNN()
        # 恢复网络权重
        self.cnn.model.load_weights(latest).expect_partial()

    def predict(self, img):
        # 以黑白方式读取图片
        # img = Image.open(image_path).convert('L')
        img_resized = cv2.resize(img, (28, 28))
        flatten_img = np.reshape(img_resized, (28, 28, 1))
        x = np.array([flatten_img / 255.0])

        # API refer: https://keras.io/models/model/
        y = self.cnn.model.predict(x)

        # 因为x只传入了一张图片，取y[0]即可
        # np.argmax()取得最大值的下标，即代表的数字
        # print(image_path)
        # print(y[0])
        # print('        -> Predict digit', np.argmax(y[0]))
        return np.argmax(y[0])


if __name__ == "__main__":
    app = Predict()
    app.predict(BASE_DIR + 'output/num0.png')
    # app.predict('num1.png')
    # app.predict('num2.png')
